import threading

class RWLock:
    def __init__(self):
        self.read_lock = threading.Lock()  # 保护读计数器的锁
        self.write_lock = threading.Lock()  # 写操作的独占锁
        self.read_count = 0  # 读操作计数器

    # 获取读锁
    def acquire_read(self):
        with self.read_lock:
            self.read_count += 1
            # 第一个读操作需要获取写锁（阻止写操作）
            if self.read_count == 1:
                self.write_lock.acquire()

    # 释放读锁
    def release_read(self):
        with self.read_lock:
            self.read_count -= 1
            # 最后一个读操作释放写锁（允许写操作）
            if self.read_count == 0:
                self.write_lock.release()

    # 获取写锁（独占）
    def acquire_write(self):
        self.write_lock.acquire()

    # 释放写锁
    def release_write(self):
        self.write_lock.release()

    # 支持 with 语句（读模式）
    def __enter__(self):
        self.acquire_read()
        return self

    # 支持 with 语句（读模式）
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.release_read()

    # 写模式的上下文管理器
    class WriteLock:
        def __init__(self, rw_lock):
            self.rw_lock = rw_lock

        def __enter__(self):
            self.rw_lock.acquire_write()
            return self

        def __exit__(self, exc_type, exc_val, exc_tb):
            self.rw_lock.release_write()

    def write_mode(self):
        return self.WriteLock(self)


# 使用示例
if __name__ == "__main__":
    rw_lock = RWLock()
    data = 0

    # 读操作函数
    def reader(id):
        for _ in range(3):
            with rw_lock:  # 读模式
                print(f"读者 {id} 读取数据: {data}")
                time.sleep(0.5)

    # 写操作函数
    def writer(id):
        for _ in range(2):
            with rw_lock.write_mode():  # 写模式
                global data
                data += 1
                print(f"写者 {id} 修改数据为: {data}")
                time.sleep(1)

    import time
    # 创建3个读者和2个写者
    readers = [threading.Thread(target=reader, args=(i,)) for i in range(3)]
    writers = [threading.Thread(target=writer, args=(i,)) for i in range(2)]

    for t in readers + writers:
        t.start()
    for t in readers + writers:
        t.join()
